from vanna.chromadb import ChromaDB_VectorStore
from vanna.ollama import Ollama
from vanna.exceptions import DependencyError, ImproperlyConfigured, ValidationError
import os
import pandas as pd
from sqlalchemy import create_engine, text
from sqlalchemy.exc import SQLAlchemyError
from typing import Union
from urllib.parse import quote_plus

class MyVanna(ChromaDB_VectorStore, Ollama):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        Ollama.__init__(self, config=config)

    def connect_to_iris(
        self,
        host: str = None,
        dbname: str = None,
        user: str = None,
        password: str = None,
        port: int = None,
        **kwargs
    ):
        try:
            import sqlalchemy_iris  # 确保 sqlalchemy-iris 已安装
        except ImportError:
            raise DependencyError(
                "You need to install required dependencies to execute this method, "
                "run command: \npip install sqlalchemy-iris"
            )

        host = host or os.getenv("HOST")
        if not host:
            raise ImproperlyConfigured("Please set your IRIS host")

        dbname = dbname or os.getenv("DATABASE")
        if not dbname:
            raise ImproperlyConfigured("Please set your IRIS database")

        user = user or os.getenv("USER")
        if not user:
            raise ImproperlyConfigured("Please set your IRIS user")

        password = password or os.getenv("PASSWORD")
        if not password:
            raise ImproperlyConfigured("Please set your IRIS password")
        password = quote_plus(password)

        port = port or os.getenv("PORT")
        if not port:
            raise ImproperlyConfigured("Please set your IRIS port")

        self.engine = None

        try:
            self.engine = create_engine(
                f"iris://{user}:{password}@{host}:{port}/{dbname}",
                **kwargs
            )
        except SQLAlchemyError as e:
            raise ValidationError(e)

        def run_sql_iris(sql: str) -> Union[pd.DataFrame, None]:
            if self.engine:
                try:
                    with self.engine.connect() as connection:
                        result = connection.execute(text(sql))
                        results = result.fetchall()
                        
                        df = pd.DataFrame(results, columns=result.keys())
                        return df
                except SQLAlchemyError as e:
                    raise ValidationError(e)
                except Exception as e:
                    raise e

        self.dialect = "T-SQL / Microsoft SQL Server"
        self.run_sql_is_set = True
        self.run_sql = run_sql_iris